#include "head.hpp"

#define STACK_SIZE 1024*8

class Schedule;
class Cor;

enum CoState {
    DEAD,
    READY,
    RUNNING,
    SUSPEND
};

using Func = std::function<void(Schedule*, void*)>;

class Schedule
{
public:
    vector<shared_ptr<Cor>> cors;
    queue<int> back_cors_que;
    int run_cor_id;

    void mainfunc();
    int create_cor(Func func, void* args);
    bool resume_cor(int cor_id);
    bool yield_cor();
    bool all_finished();

    Schedule() {
        run_cor_id = -1;
    }
};

class Cor
{
public:
    ucontext_t context;
    ucontext_t pre_context;
    Func func;
    void* args;
    enum CoState state;
    char stack_[STACK_SIZE];

    void set_run_env(Func func, void* args, enum CoState state) {
        this->func = func;
        this->args = args;
        this->state = state;
    }
};

void Schedule::mainfunc() {
    // cout << "mainfunc" << endl;
    assert(run_cor_id >= 0 && run_cor_id < cors.size());
    auto cor = cors[run_cor_id];
    cor->func(this, cor->args);
    cor->state = DEAD;
    back_cors_que.push(run_cor_id);
}

int Schedule::create_cor(Func func, void* args) {
    if (!back_cors_que.empty()) {
        int id = back_cors_que.front();
        back_cors_que.pop();
        assert(id >= 0 && id < cors.size());
        assert(cors[id]->state == DEAD);

        auto cor = cors[id];
        makecontext(&(cor->context), reinterpret_cast<void(*)(void)>(&Schedule::mainfunc), 1, this);
        cor->set_run_env(func, args, READY);
        return id;
    } else {
        auto cor = make_shared<Cor>();
        cor->set_run_env(func, args, READY);
        getcontext(&(cor->context));
        cor->context.uc_link = &(cor->pre_context);
        cor->context.uc_stack.ss_flags = 0;
        cor->context.uc_stack.ss_size = STACK_SIZE;
        cor->context.uc_stack.ss_sp = cor->stack_;
        makecontext(&(cor->context), reinterpret_cast<void(*)(void)>(&Schedule::mainfunc), 1, this);
        cors.push_back(cor);
        return cors.size()-1;
    }
}

bool Schedule::resume_cor(int cor_id) {
    // cout << "resume" << cor_id << endl;
    int pre_cor_id = run_cor_id;
    assert(cor_id >= 0 && cor_id < cors.size());
    auto cor = cors[cor_id];
    if (cor->state == READY || cor->state == SUSPEND) {
        run_cor_id = cor_id;
        cor->state = RUNNING;
        swapcontext(&(cor->pre_context), &(cor->context));
        run_cor_id = pre_cor_id;
    }
    else return false;
    return true;
}

bool Schedule::yield_cor() {
    // cout << "yiled: " << run_cor_id << endl;
    assert(run_cor_id >= 0 && run_cor_id < cors.size());
    auto cor = cors[run_cor_id];
    cor->state = SUSPEND;
    // don't need worry about run_cor_id, because resume_cor deals with it
    swapcontext(&(cor->context), &(cor->pre_context));
}

bool Schedule::all_finished() {
    return back_cors_que.size() == cors.size();
}





